/***************************************************************************
 *   Copyright (C) 2010 by Oleg Goncharov  *
 *   $EMAIL$                           *                          
 *                                                                         *
 *   This file is part of ChessVision.                                     *
 *                                                                         *
 *   ChessVision is free software; you can redistribute it and/or modify   *
 *   it under the terms of the GNU General Public License as published by  *
 *   the Free Software Foundation; either version 2 of the License, or     *
 *   (at your option) any later version.                                   *
 *                                                                         *
 *   This program is distributed in the hope that it will be useful,       *
 *   but WITHOUT ANY WARRANTY; without even the implied warranty of        *
 *   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the         *
 *   GNU General Public License for more details.                          *
 *                                                                         *
 *   You should have received a copy of the GNU General Public License     *
 *   along with this program; if not, write to the                         *
 *   Free Software Foundation, Inc.,                                       *
 *   59 Temple Place - Suite 330, Boston, MA  02111-1307, USA.             *
 ***************************************************************************/
#include <opencv/cxcore.h>
#include <opencv/cv.h>
#include <opencv/highgui.h>
#include <iostream>
#include <cmath>

#include "vision/houghtransform.h"
#include "chtransdetector.h"

#define HSTEP_THEREHOLD 3
#define TETA_THEREHOLD 3

namespace Chess {

CHTransDetector::CHTransDetector() {
	rho_step = 0.7;
	teta_step = 0.03f;
	tweak_rho = 100.0f;
	tweak_teta = 0.03f;
	therehold1 = 30;
	therehold2 = 90;
	min_board_size = 0.5;
}


CHTransDetector::~CHTransDetector()
{
}
//FIXME Do not work with ideal.png
bool CHTransDetector::DetectChessboard(IplImage * img, CBoardCells& cells) { 
	CvSize size = cvGetSize(img);
	CvSize hsize = cvCalcHoughSize(size, rho_step, teta_step);
	
	IplImage * grey = cvCreateImage(size, IPL_DEPTH_8U, 1);
	IplImage * edge = cvCreateImage(size, IPL_DEPTH_8U, 1);
	IplImage * ht = cvCreateImage(hsize , IPL_DEPTH_32S, 1);
	
	cvCvtColor(img, grey, CV_BGR2GRAY);
	cvCanny(grey, edge, therehold1, therehold2);
	//cvSobel(grey, edge, 1, 1, 5);
	cvHoughTransform(edge, ht);
	

	/*cvNamedWindow("edge", 0);
	cvShowImage("edge",edge);
	
	cvNamedWindow("ht", 0);
	cvShowImage("ht", ht);
	cvWaitKey(0);*/
	
	cvReleaseImage(&edge);
	cvReleaseImage(&grey);
	
	SBoardPosition bp;
	
	int half_hheight = hsize.height/2;
	int half_hwidth = hsize.width/2;
	
	float R = sqrt(hsize.height*hsize.height + hsize.width*hsize.width);
	float alpha = atanf(float(size.height)/float(size.width));
	float teta = 0;
	
	int start_h, stop_h;
	int start_rho, stop_rho;
	//register int index;
	
	int max_sum = 0, sum;
	int best_i_rho1, best_i_rho2;
	int best_i_h1, best_i_h2;
	int best_i_teta1, best_i_teta2;
	
	start_h = int( min_board_size * MIN(size.width, size.height) / 8 / rho_step );
	
	CvMat matheader;
	CvMat * htmat = cvGetMat(ht, &matheader);
	int teta_shift = htmat->step;
	int htelem_size = CV_ELEM_SIZE(htmat->type);
	int h_shift;
	uchar * ht_row_ptr = htmat->data.ptr;
	uchar * ht_ptr;
	
	start_rho = half_hwidth;
	for(int i_teta = 0; i_teta < half_hheight; i_teta++, teta += teta_step, ht_row_ptr += teta_shift) {
		stop_rho = int(half_hwidth * cosf(teta - alpha));
		stop_h = stop_rho/8;
		stop_rho += half_hwidth;
		
		for(int i_h = start_h; i_h <= stop_h; i_h++) {
			h_shift = i_h*htelem_size;
			for(int i_rho = stop_rho - 8*i_h; i_rho >= start_rho; i_rho--) {
				ht_ptr = ht_row_ptr + htelem_size*i_rho;
				//index = i_rho;
				sum = 0;
				//for(int i = 0; i < 9; i++, index += i_h) sum += cvGetReal2D(ht, i_teta, index );
				for(int i = 0; i < 9; i++, ht_ptr += h_shift) sum += *((int*) ht_ptr);
				if (sum > max_sum) {
					max_sum = sum;
					best_i_rho1 = i_rho;
					best_i_h1 = i_h;
					best_i_teta1 = i_teta; 
				}
			}
		}
	}
	
	max_sum = 0;
	for(int i_teta = half_hheight; i_teta < hsize.height; i_teta++, teta += teta_step, ht_row_ptr += teta_shift) {
		start_rho = half_hwidth + int(size.width*cosf(teta)/rho_step);
		stop_rho = half_hwidth + int(size.height*sinf(teta)/rho_step);
		stop_h = (stop_rho - start_rho)/8;
		
		for(int i_h = start_h; i_h <= stop_h; i_h++) {
			h_shift = i_h*htelem_size;
			for(int i_rho = stop_rho - 8*i_h; i_rho >= start_rho; i_rho--) {
				ht_ptr = ht_row_ptr + htelem_size*i_rho;
				//index = i_rho;
				sum = 0;
				//for(int i = 0; i < 9; i++, index += i_h) sum += cvGetReal2D(ht, i_teta, index );
				for(int i = 0; i < 9; i++, ht_ptr += h_shift) sum += *((int*) ht_ptr);
				if (sum > max_sum) {
					max_sum = sum;
					best_i_rho2 = i_rho;
					best_i_h2 = i_h;
					best_i_teta2 = i_teta; 
				}
			}
		}
	}
	
	
	/*std::cout << "h1 = " << best_i_h1 * rho_step << std::endl;
	std::cout << "h2 = " << best_i_h2 * rho_step << std::endl;
	std::cout << "teta1 = " << best_teta1  << std::endl;
	std::cout << "teta2 = " << best_teta2  << std::endl;
	
	IplImage * ht_col = cvCreateImage(hsize , IPL_DEPTH_16U, 3);
	IplImage * ht16 = cvCreateImage(hsize , IPL_DEPTH_16U, 1);
	cvConvert(ht, ht16);
	cvCvtColor(ht16, ht_col, CV_GRAY2BGR);
	index = best_i_rho1;
	int i_teta = best_teta1 / teta_step;
	for(int i = 0; i < 9; i++, index += best_i_h1) 
		cvCircle( ht_col, cvPoint(index, i_teta), 3, cvScalar(0,0,65000), -1);
	index = best_i_rho2;
	i_teta = best_teta2 / teta_step;
	for(int i = 0; i < 9; i++, index += best_i_h2) 
		cvCircle( ht_col, cvPoint(index, i_teta), 3, cvScalar(0,0,65000), -1);
	
	cvShowImage("ht", ht_col);
	cvWaitKey(0);
	*/
	
//TODO int -> cvRound
	int i_drho = 2*int(tweak_rho / rho_step) + 1;
	int i_dteta = 2*int(tweak_teta / teta_step) + 1;	
	//std::cout << "delta_h = " << best_i_h1 - best_i_h2 << " delta_teta = " << (best_i_teta2 - best_i_teta1)*teta_step - M_PI/2 << std::endl;
	bool looksRectangular = abs(best_i_h1 - best_i_h2) <= HSTEP_THEREHOLD && abs((best_i_teta2 - best_i_teta1) - int(M_PI/(2*teta_step))) < 4*i_dteta;
	if (looksRectangular) {
		bp.h = (best_i_h1 + best_i_h2) * rho_step / 2;
		bp.teta = ((best_i_teta1 + best_i_teta2)*teta_step - M_PI/2) / 2;
		float rho1 = (best_i_rho1 + 4*best_i_h1 - half_hwidth)*rho_step;
		float rho2 = (best_i_rho2 + 4*best_i_h1 - half_hwidth)*rho_step;
		bp.x = rho1*cos(bp.teta) - rho2*sin(bp.teta);
		bp.y = rho1*sin(bp.teta) + rho2*cos(bp.teta);
		
		//if (Integral(grey, bp, int(bp.h/2)) < 0) bp.teta += M_PI/2.0f;
		//cells.Init(bp.x, bp.y, bp.teta, best_i_h1*rho_step, best_i_h2*rho_step);
		
		int i_rho1 = best_i_rho1;
		int i_rho2 = best_i_rho2;
		cv::Size rsize1(MIN(i_drho, best_i_h1/4), i_dteta);
		cv::Size rsize2(MIN(i_drho, best_i_h2/4), i_dteta);
		Line2f lines1[9];
		Line2f lines2[9];
		
		for(int i = 0; i < 9; i++) {
			lines1[i] = RefineLine(ht, cv::Point(i_rho1, best_i_teta1) , rsize1); 
			lines2[i] = RefineLine(ht, cv::Point(i_rho2, best_i_teta2) , rsize2); 
			i_rho1 += best_i_h1;
			i_rho2 += best_i_h2;
		}
		cells.Init(lines1, lines2);
		//cells.Init(bp);
	}
	
	cvReleaseImage(&ht);
	
	return looksRectangular;
}

Line2f CHTransDetector::RefineLine(const cv::Mat& ht, const cv::Point& pt, const cv::Size& rect_size) {
	cv::Rect rect(pt.x - rect_size.width/2, pt.y - rect_size.height/2, rect_size.width, rect_size.height);
	rect &= cv::Rect(cv::Point(0, 0), ht.size());
	
	cv::Mat region(ht, rect);
	cv::Point max_loc;
	cv::minMaxLoc(region, 0, 0, 0, &max_loc);
	max_loc += rect.tl();
	
	return Line2f((max_loc.x - ht.cols/2) * rho_step, max_loc.y * teta_step);
}

}

